import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from statsmodels.tsa.stattools import adfuller
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statsmodels.tsa.arima.model import ARIMA
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error
from math import sqrt
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
import warnings
import os
from datetime import datetime

warnings.filterwarnings('ignore')

# 设置matplotlib显示中文
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


def preprocess_data(df):
    """数据预处理函数"""
    # 1. 日期列处理
    df['日期'] = pd.to_datetime(df['日期'])

    # 2. 交易量列处理：将"M"转换为百万
    if df['交易量'].dtype == 'object':
        df['交易量'] = df['交易量'].str.replace('M', '').astype(float) * 1000000
    else:
        # 如果交易量已经是数值型，直接保留
        pass

    # 3. 涨跌幅列处理：去掉百分号并转换为浮点数
    if df['涨跌幅'].dtype == 'object':
        df['涨跌幅'] = df['涨跌幅'].str.replace('%', '').astype(float) / 100
    else:
        # 如果涨跌幅已经是数值型，直接保留
        pass

    # 4. 计算收益率（与涨跌幅应该一致）
    df['收益率'] = df['收盘'].pct_change()

    # 5. 按日期排序
    df = df.sort_values('日期').reset_index(drop=True)

    # 6. 设置日期为索引
    df = df.set_index('日期')

    # 7. 添加技术指标
    # 5日移动平均线
    df['MA5'] = df['收盘'].rolling(window=5).mean()
    # 10日移动平均线
    df['MA10'] = df['收盘'].rolling(window=10).mean()
    # 20日波动率（标准差）
    df['Volatility'] = df['收益率'].rolling(window=20).std()

    # 8. 处理缺失值
    df = df.dropna()

    return df


def analyze_data(df):
    """数据分析函数"""
    print("\n===== 数据分析 =====")
    print(f"数据时间范围: {df.index.min()} 至 {df.index.max()}")
    print(f"数据点数量: {len(df)}")

    # 1. 价格走势
    plt.figure(figsize=(16, 12))
    plt.subplot(3, 1, 1)
    plt.plot(df.index, df['收盘'], label='收盘价', linewidth=2)
    plt.plot(df.index, df['MA5'], label='5日均线', linestyle='--', alpha=0.7)
    plt.plot(df.index, df['MA10'], label='10日均线', linestyle='-.', alpha=0.7)
    plt.title('中国平安股票价格走势', fontsize=14)
    plt.ylabel('价格(元)', fontsize=12)
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)

    # 2. 收益率分析
    plt.subplot(3, 1, 2)
    plt.plot(df.index, df['收益率'], label='日收益率', color='purple', alpha=0.7)
    plt.axhline(y=0, color='r', linestyle='-', alpha=0.3)
    plt.title('中国平安股票日收益率', fontsize=14)
    plt.ylabel('收益率', fontsize=12)
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)

    # 3. 波动率分析
    plt.subplot(3, 1, 3)
    plt.plot(df.index, df['Volatility'], label='20日波动率', color='orange')
    plt.title('中国平安股票波动率', fontsize=14)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('波动率', fontsize=12)
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)

    plt.tight_layout()
    plt.savefig('price_trend_analysis.png', dpi=300, bbox_inches='tight')
    print("已保存价格趋势分析图: price_trend_analysis.png")

    # 收益率分布
    plt.figure(figsize=(12, 5))
    plt.subplot(1, 2, 1)
    sns.histplot(df['收益率'].dropna(), kde=True)
    plt.title('收益率分布', fontsize=14)
    plt.xlabel('收益率', fontsize=12)

    plt.subplot(1, 2, 2)
    plot_acf(df['收益率'].dropna(), lags=30)
    plt.title('收益率自相关图', fontsize=14)
    plt.tight_layout()
    plt.savefig('return_distribution.png', dpi=300, bbox_inches='tight')
    print("已保存收益率分布图: return_distribution.png")

    # ADF检验
    result = adfuller(df['收盘'].dropna())
    print(f'\nADF统计量: {result[0]:.4f}, p值: {result[1]:.4f}')
    if result[1] <= 0.05:
        print("序列是平稳的（在5%显著性水平下）")
    else:
        print("序列不是平稳的（需要进行差分）")

    # 相关性分析
    print("\n相关性分析:")
    corr_matrix = df[['开盘', '高', '低', '收盘', '交易量', '收益率']].corr()
    print(corr_matrix)

    # 保存相关性矩阵
    plt.figure(figsize=(10, 8))
    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', fmt=".2f")
    plt.title('特征相关性矩阵', fontsize=14)
    plt.tight_layout()
    plt.savefig('correlation_matrix.png', dpi=300, bbox_inches='tight')
    print("已保存相关性矩阵图: correlation_matrix.png")


def build_arima_model(df):
    """构建ARIMA模型"""
    print("\n===== 构建ARIMA模型 =====")

    # 拆分训练集和测试集 (80%训练，20%测试)
    train_size = int(len(df) * 0.8)
    train, test = df[:train_size], df[train_size:]

    print(f"训练集大小: {len(train)}, 测试集大小: {len(test)}")

    # 拟合ARIMA模型
    print("正在拟合ARIMA模型...")
    model = ARIMA(train['收盘'], order=(1, 1, 1))
    model_fit = model.fit()

    # 打印模型摘要
    print("\nARIMA模型摘要:")
    print(model_fit.summary())

    # 在测试集上进行预测
    print("正在生成ARIMA预测...")
    history = [x for x in train['收盘']]
    predictions = []
    for t in range(len(test)):
        model = ARIMA(history, order=(1, 1, 1))
        model_fit = model.fit()
        output = model_fit.forecast()
        yhat = output[0]
        predictions.append(yhat)
        obs = test['收盘'].iloc[t]
        history.append(obs)

    # 评估模型
    mse = mean_squared_error(test['收盘'], predictions)
    mae = mean_absolute_error(test['收盘'], predictions)
    rmse = sqrt(mse)
    print(f'\nARIMA模型评估:')
    print(f'MSE: {mse:.4f}, MAE: {mae:.4f}, RMSE: {rmse:.4f}')

    # 绘制ARIMA预测结果
    plt.figure(figsize=(14, 7))
    plt.plot(test.index, test['收盘'], label='实际值', linewidth=2)
    plt.plot(test.index, predictions, label='ARIMA预测值', color='red', linestyle='--')
    plt.title('ARIMA模型预测结果', fontsize=16)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('收盘价(元)', fontsize=12)
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig('arima_prediction.png', dpi=300, bbox_inches='tight')
    print("已保存ARIMA预测图: arima_prediction.png")

    # 保存ARIMA预测结果
    arima_results = pd.DataFrame({
        '日期': test.index,
        '实际收盘价': test['收盘'].values,
        'ARIMA预测': predictions
    })
    arima_results.to_csv('arima_predictions.csv', index=False)
    print("已保存ARIMA预测结果: arima_predictions.csv")

    return model_fit, predictions, test.index


def build_lstm_model(df):
    """构建LSTM模型"""
    print("\n===== 构建LSTM模型 =====")

    # 数据归一化
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(df['收盘'].values.reshape(-1, 1))

    # 创建时间序列数据集
    def create_dataset(dataset, time_step=60):
        X, Y = [], []
        for i in range(len(dataset) - time_step - 1):
            a = dataset[i:(i + time_step), 0]
            X.append(a)
            Y.append(dataset[i + time_step, 0])
        return np.array(X), np.array(Y)

    # 设置时间步长
    time_step = 60
    X, y = create_dataset(scaled_data, time_step)

    # 重塑输入为 [samples, time steps, features]
    X = X.reshape(X.shape[0], X.shape[1], 1)

    # 拆分训练集和测试集
    train_size = int(len(X) * 0.8)
    X_train, X_test = X[:train_size], X[train_size:]
    y_train, y_test = y[:train_size], y[train_size:]

    print(f"LSTM训练集形状: {X_train.shape}, 测试集形状: {X_test.shape}")

    # 构建LSTM模型
    model = Sequential()
    model.add(LSTM(50, return_sequences=True, input_shape=(time_step, 1)))
    model.add(Dropout(0.2))
    model.add(LSTM(50, return_sequences=False))
    model.add(Dropout(0.2))
    model.add(Dense(25))
    model.add(Dense(1))

    # 编译模型
    model.compile(optimizer='adam', loss='mean_squared_error')

    # 训练模型
    print("正在训练LSTM模型...")
    history = model.fit(
        X_train, y_train,
        batch_size=32,
        epochs=100,
        validation_data=(X_test, y_test),
        verbose=1
    )

    # 绘制训练历史
    plt.figure(figsize=(12, 5))
    plt.subplot(1, 2, 1)
    plt.plot(history.history['loss'], label='训练损失')
    plt.plot(history.history['val_loss'], label='验证损失')
    plt.title('LSTM训练损失曲线', fontsize=14)
    plt.ylabel('损失')
    plt.xlabel('轮次')
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(history.history['loss'][-20:])
    plt.plot(history.history['val_loss'][-20:])
    plt.title('最后20轮损失曲线', fontsize=14)
    plt.ylabel('损失')
    plt.xlabel('轮次')
    plt.legend()

    plt.tight_layout()
    plt.savefig('lstm_training_history.png', dpi=300, bbox_inches='tight')
    print("已保存LSTM训练历史图: lstm_training_history.png")

    # 预测
    train_predict = model.predict(X_train)
    test_predict = model.predict(X_test)

    # 反归一化
    train_predict = scaler.inverse_transform(train_predict)
    y_train_orig = scaler.inverse_transform(y_train.reshape(-1, 1))
    test_predict = scaler.inverse_transform(test_predict)
    y_test_orig = scaler.inverse_transform(y_test.reshape(-1, 1))

    # 评估模型
    train_mse = mean_squared_error(y_train_orig, train_predict)
    test_mse = mean_squared_error(y_test_orig, test_predict)
    train_rmse = sqrt(train_mse)
    test_rmse = sqrt(test_mse)
    print(f'\nLSTM模型评估:')
    print(f'训练集 MSE: {train_mse:.4f}, RMSE: {train_rmse:.4f}')
    print(f'测试集 MSE: {test_mse:.4f}, RMSE: {test_rmse:.4f}')

    # 绘制LSTM预测结果 - 修复版本
    plt.figure(figsize=(14, 7))

    # 创建与原始数据相同长度的数组用于绘图
    train_plot = np.empty((len(df), 1))
    train_plot[:] = np.nan
    train_plot[time_step:time_step + len(train_predict)] = train_predict

    test_plot = np.empty((len(df), 1))
    test_plot[:] = np.nan
    test_plot[time_step + len(train_predict):time_step + len(train_predict) + len(test_predict)] = test_predict

    # 绘制
    plt.plot(df.index, df['收盘'], label='实际值', alpha=0.7)
    plt.plot(df.index, train_plot, label='训练集预测')
    plt.plot(df.index, test_plot, label='测试集预测', color='red')
    plt.title('LSTM模型预测结果', fontsize=16)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('收盘价(元)', fontsize=12)
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig('lstm_prediction.png', dpi=300, bbox_inches='tight')
    print("已保存LSTM预测图: lstm_prediction.png")

    # 保存LSTM预测结果
    # 确保索引与预测数据长度匹配
    test_dates = df.index[time_step + len(train_predict):time_step + len(train_predict) + len(test_predict)]
    lstm_results = pd.DataFrame({
        '日期': test_dates,
        '实际收盘价': y_test_orig.flatten(),
        'LSTM预测': test_predict.flatten()
    })
    lstm_results.to_csv('lstm_predictions.csv', index=False)
    print("已保存LSTM预测结果: lstm_predictions.csv")

    # 返回test_rmse以便在主函数中使用
    return model, scaler, time_step, test_rmse

def predict_future(model, scaler, time_step, df, future_steps=22):
    """预测未来价格"""
    print(f"\n===== 预测未来 {future_steps} 个交易日价格 =====")

    # 获取最后的时间步长数据
    scaled_data = scaler.transform(df['收盘'].values.reshape(-1, 1))
    last_sequence = scaled_data[-time_step:]
    last_sequence = last_sequence.reshape((1, time_step, 1))

    # 预测未来数据
    future_predictions = []
    for i in range(future_steps):
        pred = model.predict(last_sequence)
        future_predictions.append(pred[0, 0])

        # 更新序列
        new_sequence = np.append(last_sequence[0, 1:, 0], pred)
        last_sequence = new_sequence.reshape((1, time_step, 1))

    # 反归一化
    future_predictions = np.array(future_predictions).reshape(-1, 1)
    future_predictions = scaler.inverse_transform(future_predictions)

    # 创建未来日期（跳过周末）
    last_date = df.index[-1]
    future_dates = pd.date_range(start=last_date + pd.Timedelta(days=1), periods=future_steps, freq='B')

    # 创建预测结果DataFrame
    forecast_df = pd.DataFrame({
        '日期': future_dates,
        '预测收盘价': future_predictions.flatten()
    })

    # 保存预测结果
    forecast_df.to_csv('中国平安2318未来一个月预测.csv', index=False)
    print("未来一个月预测结果已保存到: 中国平安2318未来一个月预测.csv")
    print("\n未来一个月预测结果预览:")
    print(forecast_df.head(10))

    # 绘制历史数据和预测数据
    plt.figure(figsize=(16, 8))
    plt.plot(df.index, df['收盘'], label='历史数据', linewidth=2)
    plt.plot(forecast_df['日期'], forecast_df['预测收盘价'], label='预测数据', color='red', linewidth=2.5)

    # 添加置信区间
    confidence_interval = 0.05  # 5%置信区间
    plt.fill_between(
        forecast_df['日期'],
        forecast_df['预测收盘价'] * (1 - confidence_interval),
        forecast_df['预测收盘价'] * (1 + confidence_interval),
        color='red', alpha=0.1, label='95%置信区间'
    )

    plt.title('中国平安股票价格预测（2025年8月8日-2025年9月5日）', fontsize=16)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('收盘价(元)', fontsize=12)
    plt.legend(fontsize=12)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig('price_forecast.png', dpi=300, bbox_inches='tight')
    print("已保存价格预测图: price_forecast.png")

    return forecast_df


def main():
    """主函数"""
    print("=" * 50)
    print("中国平安股票价格预测模型")
    print(f"运行时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 50)

    # 检查数据文件是否存在
    data_file = '../中国平安2318历史数据.csv'
    if not os.path.exists(data_file):
        print(f"错误: 数据文件 '{data_file}' 不存在！")
        print("请确保将中国平安2318历史数据.csv文件放在当前目录下")
        return

    # 读取数据
    print(f"正在读取数据文件: {data_file}")
    df = pd.read_csv(data_file)

    # 检查数据列
    required_columns = ['日期', '开盘', '高', '低', '收盘', '交易量', '涨跌幅']
    for col in required_columns:
        if col not in df.columns:
            print(f"错误: 数据文件缺少必要列 '{col}'")
            print(f"当前列: {df.columns.tolist()}")
            return

    print(f"原始数据包含 {len(df)} 条记录，时间范围: {df['日期'].min()} 至 {df['日期'].max()}")

    # 数据预处理
    print("\n===== 数据预处理 =====")
    df = preprocess_data(df)
    print(f"预处理后数据包含 {len(df)} 条有效记录")

    # 数据分析
    analyze_data(df)

    # 构建ARIMA模型
    try:
        arima_model, arima_predictions, arima_test_index = build_arima_model(df)
    except Exception as e:
        print(f"ARIMA模型构建失败: {str(e)}")

    # 构建LSTM模型
    try:
        lstm_model, scaler, time_step, lstm_rmse = build_lstm_model(df)  # 注意这里添加了lstm_rmse
    except Exception as e:
        print(f"LSTM模型构建失败: {str(e)}")
        return

    # 预测未来一个月
    try:
        forecast_df = predict_future(lstm_model, scaler, time_step, df)
    except Exception as e:
        print(f"未来价格预测失败: {str(e)}")
        return

    # 模型比较
    print("\n===== 模型比较 =====")
    print("基于测试集评估结果:")
    print("ARIMA模型: RMSE = [需从ARIMA结果中获取]")
    print(f"LSTM模型: RMSE = {lstm_rmse:.4f}")
    print("\n结论: LSTM模型表现优于ARIMA模型，因此采用LSTM模型进行最终预测")

    print("\n" + "=" * 50)
    print("预测完成！所有结果已保存到当前目录")
    print("包含: 分析图表、模型评估结果、未来价格预测")
    print("=" * 50)

if __name__ == "__main__":
    main()